

<!DOCTYPE html>


<html lang="en" data-content_root="" >

  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>brevitas.ops.autograd_ste_ops &#8212; Brevitas Documentation - dev</title>
  
  
  
  <script data-cfasync="false">
    document.documentElement.dataset.mode = localStorage.getItem("mode") || "";
    document.documentElement.dataset.theme = localStorage.getItem("theme") || "";
  </script>
  
  <!-- Loaded before other Sphinx assets -->
  <link href="../../../_static/styles/theme.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
<link href="../../../_static/styles/bootstrap.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
<link href="../../../_static/styles/pydata-sphinx-theme.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />

  
  <link href="../../../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
  <link rel="preload" as="font" type="font/woff2" crossorigin href="../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" />

    <link rel="stylesheet" type="text/css" href="../../../_static/pygments.css" />
    <link rel="stylesheet" type="text/css" href="../../../_static/sg_gallery.css" />
  
  <!-- Pre-loaded scripts that we'll load fully later -->
  <link rel="preload" as="script" href="../../../_static/scripts/bootstrap.js?digest=3ee479438cf8b5e0d341" />
<link rel="preload" as="script" href="../../../_static/scripts/pydata-sphinx-theme.js?digest=3ee479438cf8b5e0d341" />
  <script src="../../../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=3ee479438cf8b5e0d341"></script>

    <script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
    <script src="../../../_static/jquery.js"></script>
    <script src="../../../_static/underscore.js"></script>
    <script src="../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
    <script src="../../../_static/doctools.js"></script>
    <script src="../../../_static/sphinx_highlight.js"></script>
    <script crossorigin="anonymous" integrity="sha256-Ae2Vz/4ePdIu6ZyI/5ZGsYnb+m0JlOmKPjt6XZ9JJkA=" src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js"></script>
    <script>DOCUMENTATION_OPTIONS.pagename = '_modules/brevitas/ops/autograd_ste_ops';</script>
    <script>
        DOCUMENTATION_OPTIONS.theme_version = '0.15.3';
        DOCUMENTATION_OPTIONS.theme_switcher_json_url = 'https://xilinx.github.io/brevitas/dev/_static/versions.json';
        DOCUMENTATION_OPTIONS.theme_switcher_version_match = 'dev';
        DOCUMENTATION_OPTIONS.show_version_warning_banner = false;
        </script>
    <link rel="author" title="About these documents" href="../../../about.html" />
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" />
  <meta name="viewport" content="width=device-width, initial-scale=1"/>
  <meta name="docsearch:language" content="en"/>
  </head>
  
  
  <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode="">

  
  
  <div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div>
  
  <div id="pst-scroll-pixel-helper"></div>
  
  <button type="button" class="btn rounded-pill" id="pst-back-to-top">
    <i class="fa-solid fa-arrow-up"></i>Back to top</button>

  
  <input type="checkbox"
          class="sidebar-toggle"
          id="pst-primary-sidebar-checkbox"/>
  <label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label>
  
  <input type="checkbox"
          class="sidebar-toggle"
          id="pst-secondary-sidebar-checkbox"/>
  <label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label>
  
  <div class="search-button__wrapper">
    <div class="search-button__overlay"></div>
    <div class="search-button__search-container">
<form class="bd-search d-flex align-items-center"
      action="../../../search.html"
      method="get">
  <i class="fa-solid fa-magnifying-glass"></i>
  <input type="search"
         class="form-control"
         name="q"
         id="search-input"
         placeholder="Search the docs ..."
         aria-label="Search the docs ..."
         autocomplete="off"
         autocorrect="off"
         autocapitalize="off"
         spellcheck="false"/>
  <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span>
</form></div>
  </div>

  <div class="pst-async-banner-revealer d-none">
  <aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside>
</div>

  
    <header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none">
<div class="bd-header__inner bd-page-width">
  <button class="sidebar-toggle primary-toggle" aria-label="Site navigation">
    <span class="fa-solid fa-bars"></span>
  </button>
  
  
  <div class="col-lg-3 navbar-header-items__start">
    
      <div class="navbar-item">

  

<a class="navbar-brand logo" href="../../../index.html">
  
  
  
  
  
    
    
      
    
    
    <img src="../../../_static/brevitas_logo_black.svg" class="logo__image only-light" alt="Brevitas Documentation - dev - Home"/>
    <script>document.write(`<img src="../../../_static/brevitas_logo_white.svg" class="logo__image only-dark" alt="Brevitas Documentation - dev - Home"/>`);</script>
  
  
</a></div>
    
  </div>
  
  <div class="col-lg-9 navbar-header-items">
    
    <div class="me-auto navbar-header-items__center">
      
        <div class="navbar-item">
<nav class="navbar-nav">
  <ul class="bd-navbar-elements navbar-nav">
    
<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../setup.html">
    Setup
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../getting_started.html">
    Getting Started
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../tutorials/index.html">
    Tutorials
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../settings.html">
    Settings
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../api_reference/index.html">
    API reference
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../user_guide/index.html">
    User Guides
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../faq.html">
    FAQ
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../about.html">
    About
  </a>
</li>

  </ul>
</nav></div>
      
    </div>
    
    
    <div class="navbar-header-items__end">
      
        <div class="navbar-item navbar-persistent--container">
          

 <script>
 document.write(`
   <button class="btn navbar-btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <i class="fa-solid fa-magnifying-glass"></i>
    <span class="search-button__default-text">Search</span>
    <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
   </button>
 `);
 </script>
        </div>
      
      
        <div class="navbar-item">

<script>
document.write(`
  <button class="btn btn-sm navbar-btn theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <span class="theme-switch nav-link" data-mode="light"><i class="fa-solid fa-sun fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="dark"><i class="fa-solid fa-moon fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="auto"><i class="fa-solid fa-circle-half-stroke fa-lg"></i></span>
  </button>
`);
</script></div>
      
    </div>
    
  </div>
  
  
    <div class="navbar-persistent--mobile">

 <script>
 document.write(`
   <button class="btn navbar-btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <i class="fa-solid fa-magnifying-glass"></i>
    <span class="search-button__default-text">Search</span>
    <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
   </button>
 `);
 </script>
    </div>
  

  
</div>

    </header>
  

  <div class="bd-container">
    <div class="bd-container__inner bd-page-width">
      
      
      
        
      
      <div class="bd-sidebar-primary bd-sidebar hide-on-wide">
        

  
  <div class="sidebar-header-items sidebar-primary__section">
    
    
      <div class="sidebar-header-items__center">
        
          
          
            <div class="navbar-item">
<nav class="navbar-nav">
  <ul class="bd-navbar-elements navbar-nav">
    
<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../setup.html">
    Setup
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../getting_started.html">
    Getting Started
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../tutorials/index.html">
    Tutorials
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../settings.html">
    Settings
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../api_reference/index.html">
    API reference
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../user_guide/index.html">
    User Guides
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../faq.html">
    FAQ
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../about.html">
    About
  </a>
</li>

  </ul>
</nav></div>
          
        
      </div>
    
    
    
      <div class="sidebar-header-items__end">
        
          <div class="navbar-item">

<script>
document.write(`
  <button class="btn btn-sm navbar-btn theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <span class="theme-switch nav-link" data-mode="light"><i class="fa-solid fa-sun fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="dark"><i class="fa-solid fa-moon fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="auto"><i class="fa-solid fa-circle-half-stroke fa-lg"></i></span>
  </button>
`);
</script></div>
        
      </div>
    
  </div>
  
  
  <div class="sidebar-primary-items__end sidebar-primary__section">
  </div>
  
  <div id="rtd-footer-container"></div>


      </div>
      
      <main id="main-content" class="bd-main" role="main">
        
        
          <div class="bd-content">
            <div class="bd-article-container">
              
              <div class="bd-header-article d-print-none">
<div class="header-article-items header-article__inner">
  
    <div class="header-article-items__start">
      
        <div class="header-article-item">



<nav aria-label="Breadcrumb" class="d-print-none">
  <ul class="bd-breadcrumbs">
    
    <li class="breadcrumb-item breadcrumb-home">
      <a href="../../../index.html" class="nav-link" aria-label="Home">
        <i class="fa-solid fa-home"></i>
      </a>
    </li>
    
    <li class="breadcrumb-item"><a href="../../index.html" class="nav-link">Module code</a></li>
    
    <li class="breadcrumb-item active" aria-current="page">brevitas.ops...</li>
  </ul>
</nav>
</div>
      
    </div>
  
  
</div>
</div>
              
              
              
                
<div id="searchbox"></div>
                <article class="bd-article">
                  
  <h1>Source code for brevitas.ops.autograd_ste_ops</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.</span>
<span class="c1"># SPDX-License-Identifier: BSD-3-Clause</span>

<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Implementation of various torch.autograd.Function with straight-through estimators.</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Tuple</span>

<span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">torch</span><span class="w"> </span><span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">torch.autograd</span><span class="w"> </span><span class="kn">import</span> <span class="n">Function</span>

<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">binary_sign</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">dpu_round</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">round_to_zero</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">tensor_clamp</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">tensor_clamp_</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;ScalarClampSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ScalarClampMinSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TensorClampSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;InplaceTensorClampSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;RoundToZeroSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;CeilSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;FloorSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;BinarySignSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TernarySignSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;RoundSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;AbsBinarySignGradFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;DPURoundSteFn&#39;</span><span class="p">,</span>
    <span class="s1">&#39;round_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;binary_sign_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ternary_sign_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;floor_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ceil_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;round_to_zero_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;scalar_clamp_min_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;scalar_clamp_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;tensor_clamp_ste_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;abs_binary_sign_grad_impl&#39;</span><span class="p">,</span>
    <span class="s1">&#39;dpu_round_ste_impl&#39;</span><span class="p">]</span>


<div class="viewcode-block" id="ScalarClampSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.ScalarClampSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">ScalarClampSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements ``torch.clamp`` with a straight-through gradient estimator</span>
<span class="sd">    for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to ``min_val`` and ``min_val``</span>
<span class="sd">    are always ``None``.</span>

<span class="sd">    ``ScalarClampSteFn.apply(*args)`` is first aliased to :func:`scalar_clamp_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.scalar_clamp_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.scalar_clamp_ste` and invoked when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.scalar_clamp_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
        <span class="k">return</span> <span class="n">grad_y</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Clip&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">min_val</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">max_val</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="ScalarClampMinSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.ScalarClampMinSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">ScalarClampMinSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements ``torch.clamp_min`` with a straight-through gradient estimator</span>
<span class="sd">    for the gradient of y w.r.t. to x, while the gradient of y w.r.t. to ``min_val`` is always</span>
<span class="sd">    ``None``.</span>

<span class="sd">    ``ScalarClampMinSteFn.apply(*args)`` is first aliased to :func:`scalar_clamp_min_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.scalar_clamp_min_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.scalar_clamp_min_ste` and invoked when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.scalar_clamp_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clamp_min</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">min_val</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
        <span class="k">return</span> <span class="n">grad_y</span><span class="p">,</span> <span class="kc">None</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="nb">float</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Clip&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">min_val</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="TensorClampSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.TensorClampSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">TensorClampSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`~brevitas.function.ops.tensor_clamp` with a</span>
<span class="sd">    straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y</span>
<span class="sd">    w.r.t. to min_val and max_val is always None.</span>

<span class="sd">    ``TensorClampSteFn.apply(*args)`` is first aliased to :func:`tensor_clamp_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.tensor_clamp_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.tensor_clamp` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.tensor_clamp` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">tensor_clamp</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
        <span class="k">return</span> <span class="n">grad_y</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">upper_cond</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Greater&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Where&#39;</span><span class="p">,</span> <span class="n">upper_cond</span><span class="p">,</span> <span class="n">max_val</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="n">lower_cond</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Less&#39;</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="n">min_val</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Where&#39;</span><span class="p">,</span> <span class="n">lower_cond</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="InplaceTensorClampSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.InplaceTensorClampSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">InplaceTensorClampSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`~brevitas.function.ops.tensor_clamp_` with a</span>
<span class="sd">    straight-through gradient estimator for the gradient of y w.r.t. to x, while the gradient of y</span>
<span class="sd">    w.r.t. to min_val and max_val is always None.</span>

<span class="sd">    ``InplaceTensorClampSteFn.apply(*args)`` is first aliased to</span>
<span class="sd">    :func:`tensor_clamp_ste_impl_(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.tensor_clamp_ste_impl_&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.tensor_clamp_` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.tensor_clamp_` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">tensor_clamp_</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
        <span class="k">return</span> <span class="n">grad_y</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">upper_cond</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Greater&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Where&#39;</span><span class="p">,</span> <span class="n">upper_cond</span><span class="p">,</span> <span class="n">max_val</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="n">lower_cond</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Less&#39;</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="n">min_val</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Where&#39;</span><span class="p">,</span> <span class="n">lower_cond</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="RoundToZeroSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.RoundToZeroSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">RoundToZeroSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`~brevitas.function.ops.round_to_zero` with a</span>
<span class="sd">    straight-through gradient estimator.</span>

<span class="sd">    ``RoundToZeroSteFn.apply(*args)`` is first aliased to :func:`round_to_zero_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.round_to_zero_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.round_to_zero_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.round_to_zero_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">round_to_zero</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="nb">abs</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Abs&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="n">sign</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Sign&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="n">floor</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Floor&#39;</span><span class="p">,</span> <span class="nb">abs</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Mul&#39;</span><span class="p">,</span> <span class="n">sign</span><span class="p">,</span> <span class="n">floor</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="DPURoundSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.DPURoundSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">DPURoundSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`~brevitas.function.ops.dpu_round` with a</span>
<span class="sd">    straight-through gradient estimator.</span>

<span class="sd">    ``DPURoundSteFn.apply(*args)`` is first aliased to :func:`dpu_round_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.dpu_round_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.dpu_round_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.dpu_round_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">dpu_round</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>


<div class="viewcode-block" id="CeilSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.CeilSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">CeilSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`torch.ceil` with a straight-through gradient estimator.</span>

<span class="sd">    ``CeilSteFn.apply(*args)`` is first aliased to :func:`ceil_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.ceil_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.ceil_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.ceil_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Ceil&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="FloorSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.FloorSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">FloorSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`torch.floor` with a straight-through gradient estimator.</span>

<span class="sd">    ``FloorSteFn.apply(*args)`` is first aliased to :func:`floor_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.floor_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.floor_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.floor_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Floor&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="BinarySignSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.BinarySignSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">BinarySignSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`~brevitas.function.ops.binary_sign` with a</span>
<span class="sd">    straight-through gradient estimator.</span>

<span class="sd">    ``BinarySignSteFn.apply(*args)`` is first aliased to</span>
<span class="sd">    :func:`binary_sign_ste_impl(*args)&lt;brevitas.ops.autograd_ste_ops.binary_sign_ste_impl&gt;`</span>
<span class="sd">    and then wrapped by :func:`~brevitas.function.ops_ste.binary_sign_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.binary_sign_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">binary_sign</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="c1"># requires ONNX opset &gt;= 12</span>
        <span class="n">positive_mask</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;GreaterOrEqual&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.</span><span class="p">))</span>
        <span class="n">negative_mask</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Less&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.</span><span class="p">))</span>
        <span class="n">positive_mask</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Cast&#39;</span><span class="p">,</span> <span class="n">positive_mask</span><span class="p">,</span> <span class="n">to_i</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">onnx</span><span class="o">.</span><span class="n">TensorProtoDataType</span><span class="o">.</span><span class="n">FLOAT</span><span class="p">)</span>
        <span class="n">negative_mask</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Cast&#39;</span><span class="p">,</span> <span class="n">negative_mask</span><span class="p">,</span> <span class="n">to_i</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">onnx</span><span class="o">.</span><span class="n">TensorProtoDataType</span><span class="o">.</span><span class="n">FLOAT</span><span class="p">)</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Sub&#39;</span><span class="p">,</span> <span class="n">positive_mask</span><span class="p">,</span> <span class="n">negative_mask</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="TernarySignSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.TernarySignSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">TernarySignSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`torch.sign` with a straight-through gradient estimator.</span>

<span class="sd">    ``TernarySignSteFn.apply(*args)`` is first aliased to :func:`ternary_sign_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.ternary_sign_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.ternary_sign_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.ternary_sign_ste` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Sign&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="RoundSteFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.RoundSteFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">RoundSteFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`torch.round` with a straight-through gradient</span>
<span class="sd">    estimator.</span>

<span class="sd">    ``RoundSteFn.apply(*args)`` is first aliased to :func:`round_ste_impl(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.round_ste_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.round_ste` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.round_ste` for details on the interface and examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">round</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="c1"># Requires ONNX opset &gt;= 11</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Round&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<div class="viewcode-block" id="AbsBinarySignGradFn"><a class="viewcode-back" href="../../../api_reference/brevitas.ops.html#brevitas.ops.autograd_ste_ops.AbsBinarySignGradFn">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">AbsBinarySignGradFn</span><span class="p">(</span><span class="n">Function</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Autograd function that implements :func:`torch.abs` with a binary-sign backward, in order to</span>
<span class="sd">    have subgradient 1 in 0. Compare with :func:`torch.abs`&#39; subgradient of 0 in 0.</span>

<span class="sd">    ``AbsBinarySignGradFn.apply(*args)`` is first aliased to :func:`abs_binary_sign_grad(*args)</span>
<span class="sd">    &lt;brevitas.ops.autograd_ste_ops.abs_binary_sign_grad_impl&gt;` and then wrapped by</span>
<span class="sd">    :func:`~brevitas.function.ops_ste.abs_binary_sign_grad` when env ``BREVITAS_JIT=0``.</span>
<span class="sd">    See :func:`~brevitas.function.ops_ste.abs_binary_sign_grad` for details on the interface and</span>
<span class="sd">    examples.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">ctx</span><span class="o">.</span><span class="n">save_for_backward</span><span class="p">(</span><span class="n">binary_sign</span><span class="p">(</span><span class="n">x</span><span class="p">)</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">int8</span><span class="p">))</span>  <span class="c1"># save some memory</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">backward</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">grad_y</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">binary_sign</span><span class="p">,</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">saved_tensors</span>
        <span class="k">return</span> <span class="n">binary_sign</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> <span class="o">*</span> <span class="n">grad_y</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">symbolic</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">):</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">op</span><span class="p">(</span><span class="s1">&#39;Abs&#39;</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span></div>


<span class="c1">#: Alias for :class:`RoundSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.RoundSteFn&gt;`</span>
<span class="n">round_ste_impl</span> <span class="o">=</span> <span class="n">RoundSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`BinarySignSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.BinarySignSteFn&gt;`</span>
<span class="n">binary_sign_ste_impl</span> <span class="o">=</span> <span class="n">BinarySignSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`TernarySignSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.TernarySignSteFn&gt;`</span>
<span class="n">ternary_sign_ste_impl</span> <span class="o">=</span> <span class="n">TernarySignSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`FloorSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.FloorSteFn&gt;`</span>
<span class="n">floor_ste_impl</span> <span class="o">=</span> <span class="n">FloorSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`CeilSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.CeilSteFn&gt;`</span>
<span class="n">ceil_ste_impl</span> <span class="o">=</span> <span class="n">CeilSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`RoundToZeroSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.RoundToZeroSteFn&gt;`</span>
<span class="n">round_to_zero_ste_impl</span> <span class="o">=</span> <span class="n">RoundToZeroSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`DPURoundSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.DPURoundSteFn&gt;`</span>
<span class="n">dpu_round_ste_impl</span> <span class="o">=</span> <span class="n">DPURoundSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`ScalarClampMinSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.ScalarClampMinSteFn&gt;`</span>
<span class="n">scalar_clamp_min_ste_impl</span> <span class="o">=</span> <span class="n">ScalarClampMinSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`ScalarClampSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.ScalarClampSteFn&gt;`</span>
<span class="n">scalar_clamp_ste_impl</span> <span class="o">=</span> <span class="n">ScalarClampSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`TensorClampSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.TensorClampSteFn&gt;`</span>
<span class="n">tensor_clamp_ste_impl</span> <span class="o">=</span> <span class="n">TensorClampSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`InplaceTensorClampSteFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.InplaceTensorClampSteFn&gt;`</span>
<span class="n">tensor_clamp_ste_impl_</span> <span class="o">=</span> <span class="n">InplaceTensorClampSteFn</span><span class="o">.</span><span class="n">apply</span>

<span class="c1">#: Alias for :class:`AbsBinarySignGradFn.apply(*args)</span>
<span class="c1">#: &lt;brevitas.ops.autograd_ste_ops.AbsBinarySignGradFn&gt;`</span>
<span class="n">abs_binary_sign_grad_impl</span> <span class="o">=</span> <span class="n">AbsBinarySignGradFn</span><span class="o">.</span><span class="n">apply</span>
</pre></div>

                </article>
              
              
              
              
              
                <footer class="prev-next-footer d-print-none">
                  
<div class="prev-next-area">
</div>
                </footer>
              
            </div>
            
            
              
            
          </div>
          <footer class="bd-footer-content">
            
          </footer>
        
      </main>
    </div>
  </div>
  
  <!-- Scripts loaded after <body> so the DOM is not blocked -->
  <script src="../../../_static/scripts/bootstrap.js?digest=3ee479438cf8b5e0d341"></script>
<script src="../../../_static/scripts/pydata-sphinx-theme.js?digest=3ee479438cf8b5e0d341"></script>

  <footer class="bd-footer">
<div class="bd-footer__inner bd-page-width">
  
    <div class="footer-items__start">
      
        <div class="footer-item">

  <p class="copyright">
    
      © Copyright 2025 - Advanced Micro Devices, Inc..
      <br/>
    
  </p>
</div>
      
        <div class="footer-item">

  <p class="sphinx-version">
    Created using <a href="https://www.sphinx-doc.org/">Sphinx</a> 5.3.0.
    <br/>
  </p>
</div>
      
    </div>
  
  
  
    <div class="footer-items__end">
      
        <div class="footer-item">
<script>
document.write(`
  <div class="version-switcher__container dropdown">
    <button id="pst-version-switcher-button-2"
      type="button"
      class="version-switcher__button btn btn-sm navbar-btn dropdown-toggle"
      data-bs-toggle="dropdown"
      aria-haspopup="listbox"
      aria-controls="pst-version-switcher-list-2"
      aria-label="Version switcher list"
    >
      Choose version  <!-- this text may get changed later by javascript -->
      <span class="caret"></span>
    </button>
    <div id="pst-version-switcher-list-2"
      class="version-switcher__menu dropdown-menu list-group-flush py-0"
      role="listbox" aria-labelledby="pst-version-switcher-button-2">
      <!-- dropdown will be populated by javascript on page load -->
    </div>
  </div>
`);
</script></div>
      
    </div>
  
</div>

  </footer>
  </body>
</html>